package cn.itcast.tags.ml.rs.df

import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.recommendation.{ALS, ALSModel}
import org.apache.spark.mllib.recommendation.{MatrixFactorizationModel, Rating}
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.types.{DoubleType, IntegerType, LongType, StructType}
import org.apache.spark.sql.{DataFrame, SparkSession}

/**
 * @author: xu
 * @desc: 使用电影评分数据集，调用协同过滤推荐算法ALS建立推荐模型：基于DataFrame
 *        1、 预测 用户（User）对某个电影（Product）的评分
 *        2、 为某个用户推荐10个电影（Products）
 *        3、 为某个电影推荐10个用户（Users）
 */
object SparkRmdAls {
  def main(args: Array[String]): Unit = {
    // TODO: 1. 创建SparkSession实例对象
    // 构建SparkSession实例对象，通过建造者模式创建
    val spark: SparkSession = SparkSession
      .builder()
      .appName(this.getClass.getSimpleName.stripSuffix("$"))
      .master("local[3]")
      .config("spark.sql.shuffle.partitions", "3")
      .getOrCreate()
    import spark.implicits._

    // 自定义schema信息
    val schema: StructType = new StructType()
      .add("userId", IntegerType, nullable = true)
      .add("movieId", IntegerType, nullable = true)
      .add("rating", DoubleType, nullable = true)
      .add("timestamp", LongType, nullable = true)

    // TODO: 2. 读取电影评分数据
    val rawRationsDF = spark.read
      .option("sep", "\t")
      .schema(schema)
      .csv("datas/als/ml-100k/u.data")
    // rawRationsDF.printSchema()
    // rawRationsDF.show()

    // TODO: 3. 在训练数据上使用 ALS 构建推荐模型
    val als = new ALS()
      // 迭代次数
      .setMaxIter(5)
      // 设置特征数
      .setRank(10)
      // 设置显示评分
      .setImplicitPrefs(false)
      // 设置Block的数目，默认10
      .setNumUserBlocks(4)
      .setNumItemBlocks(4)
      // 设置用户和物品列
      .setUserCol("userId")
      .setItemCol("movieId")
      .setRatingCol("rating")
    val mlAlsModel: ALSModel = als.fit(rawRationsDF)

    // TODO: 4. 获取模型中两个因子矩阵，用户和物品
    val userFeatures: DataFrame = mlAlsModel.userFactors
    // userFeatures.show(truncate = false)
    println("=========================================================")
    val itemFeatures = mlAlsModel.itemFactors
    // itemFeatures.show(truncate = false)

    // TODO: 5.推荐与预测评分
    // 5.1 预测用户对产品的评分
    val predictRating: DataFrame = mlAlsModel
      .setUserCol("userId")
      .setItemCol("movieId")
      .setPredictionCol("predictionRating")
      .transform(rawRationsDF)
    predictRating.show(truncate = false)
    println("==========================================")

    // 5.2 给用户推荐十部电影
    val recommendTenMovies: DataFrame = mlAlsModel.recommendForAllUsers(10)
    recommendTenMovies.show(truncate = false)
    println("==========================================")

    // 5.3 给电影推荐十个用户
    val recommendTenUsers: DataFrame = mlAlsModel.recommendForAllItems(10)
    recommendTenUsers.show(truncate = false)

    // TODO: 6.模型评估，使用RMSE评估模型，值越小，表示误差越小，模型越好
    val evaluator = new RegressionEvaluator()
      .setLabelCol("rating")
      .setPredictionCol("predictionRating")
      .setMetricName("rmse")
    val rmse = evaluator.evaluate(predictRating)
    println(s"Root-mean-square error = $rmse")

    // TODO: 7. 将训练的模型进行保存，以便后期使用
    mlAlsModel.save("datas/als/ml-als-model")

    // TODO: 8. 从文件系统中加载保存的模型，给用户推荐十部电影
    val loadAlsModel = ALSModel.load("datas/als/ml-als-model")
    loadAlsModel.recommendForAllUsers(10)
      .show(5, truncate = false)

    // 关闭资源
    spark.stop()
  }
}
